package com.huatai.datacenter.utils;

import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.admin.*;
import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.config.ConfigResource;
import org.apache.kafka.common.config.TopicConfig;
import org.apache.kafka.common.requests.DescribeLogDirsResponse;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ExecutionException;

/**
 * @author Lion
 * @date 2023/3/29  11:09
 */
@Slf4j
public class KafkaAdmins implements Cloneable {

	private int DEFAULT_TIME_OUT_SECOND = 60;

	private AdminClient adminClient;

	public KafkaAdmins(Properties adminProps){
		this.adminClient = KafkaAdminClient.create(adminProps);
	}

	/**
	 * @author Lion
	 * @description 获取集群配置信息
	 * @para
	 * @return org.apache.kafka.clients.admin.DescribeClusterResult
	 */
	public DescribeClusterResult describeCluster(){
		return this.adminClient.describeCluster();
	}

	/**
	 * 获取集群配置信息
	 */
	public DescribeClusterResult descCluster() {
		return this.adminClient.describeCluster();
	}

	/**
	 * @author Lion
	 * @description 列出所有kafka topic信息，通过ListTopicOptions可以设置超时和是否为kafka内部topic eg：设置超时时间
	 * 设置列出所有topic（包含kafka内部topic）ListTopicOptions options = new ListTopicOptions();options.listInternal(true); options.timeoutMs(new Integer(3000));
	 * @param  options
	 * @return java.util.Map<java.lang.String,org.apache.kafka.clients.admin.TopicListing>
	 */
	public Map<String, TopicListing> listTopics(ListTopicsOptions options) throws ExecutionException, InterruptedException {
		ListTopicsResult listTopicsResult = this.adminClient.listTopics(options);
		KafkaFuture<Map<String, TopicListing>> kafkaFuture = listTopicsResult.namesToListings();
		return kafkaFuture.get();
	}

	/**
	 * @author Lion
	 * @description 列出所有kafka topic 信息  默认不会列出kafka内部的topic 超时时间默认
	 * @param
	 * @return java.util.Set<java.lang.String>
	 */
	public Set<String> listTopics() throws ExecutionException, InterruptedException {
		ListTopicsResult listTopicsResult = this.adminClient.listTopics();
		return listTopicsResult.names().get();

	}

	/**
	 * @author Lion
	 * @description 检查topic是否存在
	 * @param topicSames
	 * @param containInternal
	 * @return java.util.Map<java.lang.String,java.lang.Boolean>
	 */
	private Map<String,Boolean> checkExists(Set<String> topicSames,boolean containInternal) throws ExecutionException, InterruptedException {
		Map<String,Boolean> result = new HashMap<>(16);
		Set<String> kafkaTopicSet = new HashSet<>(16);
		if (containInternal){
			ListTopicsOptions options = new ListTopicsOptions();
			options.listInternal(true);
			kafkaTopicSet.addAll(this.listTopics(options).keySet());
		}else {
			kafkaTopicSet.addAll(this.listTopics());
		}
		for (String topicName : topicSames){
			result.put(topicName,kafkaTopicSet.contains(topicName));
		}
		return result;
	}

	/**
	 * @author Lion
	 * @description 检查topic是否存在，默认不检查kafka内部的topic
	 * @param topicName
	 * @return boolean
	 */
	public boolean checkExists(String topicName) throws ExecutionException, InterruptedException {
		Set<String> set = new HashSet<>(16);
		set.add(topicName);
		return this.checkExists(set,false).get(topicName);
	}

	/**
	 * @author Lion
	 * @description 列出某个具体topic，详细配置信息 使用ConfigResource.Type.TOPIC  / ConfigResource.Type.BROKER 来判断资源类型
	 * @param type
	 * @param resourceName
	 * @return org.apache.kafka.clients.admin.Config
	 */
	public Config descConfigs(ConfigResource.Type type,String resourceName) throws ExecutionException, InterruptedException {
		ConfigResource resource = new ConfigResource(type,resourceName);
		DescribeConfigsResult describeClusterResult = this.adminClient.describeConfigs(Collections.singleton(resource));
		Map<ConfigResource,Config> topicConfig = describeClusterResult.all().get();
		return topicConfig.get(resource);
	}

	/**
	 * @author Lion
	 * @description 单独创建topic
	 * @param topicName
	 * @param numPartitions
	 * @param replicationFactor
	 * @return boolean
	 */
	public boolean createTopic(String topicName,int numPartitions,short replicationFactor){
		return this.createTopic(topicName,numPartitions,replicationFactor,null);
	}

	/**
	 * @author Lion
	 * @description 创建topic 指定topicName 分区数Partitions 与副本数replicationFactor
	 * @param topicName
	 * @param numPartitions
	 * @param replicationFactor
	 * @return boolean
	 */
	public boolean createTopic(String topicName,int numPartitions,short replicationFactor,Map<String, String> topicConfig){
		boolean success = false;
		NewTopic newTopic = new NewTopic(topicName,numPartitions,replicationFactor);
		newTopic.configs(topicConfig);
		CreateTopicsResult createTopicsResult = this.adminClient.createTopics(Collections.singleton(newTopic));
		try {
			createTopicsResult.values().get(topicName).isCompletedExceptionally();
			success = true;
		}catch (Exception e){
			log.error("",e);
		}
		return success;
	}

	/**
	 * @author Lion
	 * @description 根据主题名称删除主题
	 * @param topicName
	 * @return boolean
	 */
	public boolean deleteTopic(String topicName){
		Map<String,Boolean> resultMap = this.deleteTopic(Collections.singleton(topicName));
		return resultMap.get(topicName);
	}

	/**
	 * @author Lion
	 * @description 删除topic
	 * @param  topics
	 * @return java.util.Map<java.lang.String,java.lang.Boolean>
	 */
	public Map<String,Boolean> deleteTopic(Collection<String> topics){
		Map<String,Boolean> result = new HashMap<>(16);
		DeleteTopicsResult deleteTopicsResult = this.adminClient.deleteTopics(topics);
		Map<String,KafkaFuture<Void>> deleteFutures = deleteTopicsResult.values();
		for (String topicName : topics){
			boolean flag = false;

			try {
				KafkaFuture<Void> kafkaFuture = deleteFutures.get(topicName);
				kafkaFuture.get();
				flag = true;
			} catch (InterruptedException e) {
				e.printStackTrace();
			} catch (ExecutionException e) {
				e.printStackTrace();
			}
			result.put(topicName,flag);
		}
		return result;
	}

	public void close() throws IOException {
		this.adminClient.close();
	}

	public Map<String, TopicDescription> descTopics(Collection<String> topicNames) {
		Map<String, TopicDescription> result = new HashMap<String, TopicDescription>(1 << 4);
		Map<String, KafkaFuture<TopicDescription>> describeFutures = this.adminClient.describeTopics(topicNames)
			.values();
		for (String topicName : topicNames) {
			try {
				result.put(topicName, describeFutures.get(topicName).get());
			} catch (Exception e) {
				log.error("desc topic:{} error.",topicName, e);
			}
		}
		return result;
	}

	/**
	 * 删除topic
	 */
	public boolean delete(String topic) {
		Map<String, Boolean> resultMap = this.delete(Collections.singleton(topic));
		return resultMap.get(topic);
	}

	/**
	 * 删除topics
	 */
	public Map<String, Boolean> delete(Collection<String> topics) {
		Map<String, Boolean> result = new HashMap<String, Boolean>(1 << 4);
		DeleteTopicsResult delRes = this.adminClient.deleteTopics(topics);
		Map<String, KafkaFuture<Void>> deleteFutures = delRes.values();
		for (String topicName : topics) {
			boolean flag = false;
			try {
				KafkaFuture<Void> kafkaFutrue = deleteFutures.get(topicName);
				kafkaFutrue.get();
				flag = true;
			} catch (Exception e) {
				log.error("", e);
			}
			result.put(topicName, flag);
		}
		return result;
	}

	/**
	 * 若 topic不存在则创建 默认只检测非kafka内部的topic
	 *
	 * @throws Exception
	 */
	public boolean createTopicIfNotExists(String topicName, int numPartitions, short replicationFactor, long ttl)
		throws Exception {
		if (!this.checkExists(topicName)) {
			Map<String, String> configs = new HashMap<>(1 << 3);
			configs.put(TopicConfig.RETENTION_MS_CONFIG, ttl + "");
			configs.put(TopicConfig.DELETE_RETENTION_MS_CONFIG, ttl + "");
			return this.createTopic(topicName, numPartitions, replicationFactor, configs);
		}
		return true;
	}

	public Map<String,Long> getTopicDiskSizeForBroker(List<Integer> brokerIds) throws ExecutionException, InterruptedException {
		Map<String,Long>  sizeMap = new HashMap<>();
		DescribeLogDirsResult describeLogDirsResult = this.adminClient.describeLogDirs(brokerIds);
		Map<Integer,Map<String, DescribeLogDirsResponse.LogDirInfo>> tmp = describeLogDirsResult.all().get();
		tmp.forEach((k,map)->{
			map.forEach((k2,v)->{
				DescribeLogDirsResponse.LogDirInfo info = v;
				Map<TopicPartition, DescribeLogDirsResponse.ReplicaInfo> replicaInfoMap = info.replicaInfos;
				replicaInfoMap.forEach((topic,replicaInfo)->{
					long size = sizeMap.getOrDefault(topic.topic(),0L);
					size += replicaInfo.size;
					sizeMap.put(topic.topic(),size);
				});
			});
		});
		return sizeMap;
	}
}
